package com.example.tree;

import java.util.PriorityQueue;
import java.util.Comparator;
import java.util.Map;
import java.util.HashMap;

class HuffmanNode {
    char ch;
    int freq;
    HuffmanNode left;
    HuffmanNode right;

    HuffmanNode(char ch, int freq) {
        this.ch = ch;
        this.freq = freq;
    }
}

public class HuffmanTree {
    // 构建哈夫曼树
    public static HuffmanNode buildHuffmanTree(char[] chars, int[] freqs) {
        PriorityQueue<HuffmanNode> priorityQueue = new PriorityQueue<>(Comparator.comparingInt(node -> node.freq));

        // 创建叶子节点并添加到优先队列中
        for (int i = 0; i < chars.length; i++) {
            HuffmanNode node = new HuffmanNode(chars[i], freqs[i]);
            priorityQueue.add(node);
        }

        // 合并节点直到只剩一个根节点
        while (priorityQueue.size() > 1) {
            HuffmanNode left = priorityQueue.poll();
            HuffmanNode right = priorityQueue.poll();

            HuffmanNode merged = new HuffmanNode('\0', left.freq + right.freq);
            merged.left = left;
            merged.right = right;

            priorityQueue.add(merged);
        }

        return priorityQueue.poll();
    }

    // 生成哈夫曼编码表
    public static Map<Character, String> generateCodes(HuffmanNode root) {
        Map<Character, String> codes = new HashMap<>();
        generateCodes(root, "", codes);
        return codes;
    }

    private static void generateCodes(HuffmanNode root, String code, Map<Character, String> codes) {
        if (root == null) {
            return;
        }

        if (root.left == null && root.right == null) {
            codes.put(root.ch, code);
            return;
        }

        generateCodes(root.left, code + "0", codes);
        generateCodes(root.right, code + "1", codes);
    }

    // 编码
    public static String encode(String text, Map<Character, String> codes) {
        StringBuilder encoded = new StringBuilder();
        for (char ch : text.toCharArray()) {
            encoded.append(codes.get(ch));
        }
        return encoded.toString();
    }

    // 解码
    public static String decode(String encoded, HuffmanNode root) {
        StringBuilder decoded = new StringBuilder();
        HuffmanNode current = root;

        for (char bit : encoded.toCharArray()) {
            if (bit == '0') {
                current = current.left;
            } else {
                current = current.right;
            }

            if (current.left == null && current.right == null) {
                decoded.append(current.ch);
                current = root;
            }
        }

        return decoded.toString();
    }

    public static void main(String[] args) {
        char[] chars = {'a', 'b', 'c', 'd', 'e', 'f'};
        int[] freqs = {5, 9, 12, 13, 16, 45};

        // 构建哈夫曼树
        HuffmanNode huffmanTree = buildHuffmanTree(chars, freqs);

        // 生成哈夫曼编码表
        Map<Character, String> codes = generateCodes(huffmanTree);

        // 输出编码表
        System.out.println("Huffman Codes:");
        for (Map.Entry<Character, String> entry : codes.entrySet()) {
            System.out.println(entry.getKey() + ": " + entry.getValue());
        }

        // 测试编码和解码
        String text = "abacab";
        System.out.println("\nOriginal Text: " + text);

        String encoded = encode(text, codes);
        System.out.println("Encoded Text: " + encoded);

        String decoded = decode(encoded, huffmanTree);
        System.out.println("Decoded Text: " + decoded);
    }
}